# -*- coding: utf-8 -*-
# @Time : 2023/4/23 17:10
# @Author : 陈鹏飞
# @Email ： 2578925789@qq.com
# @File : readNpy
# @Description : 读取保存的模型检测匹配结果并计算AP和PR曲线

import os
import numpy as np
import json
from pathlib import Path
import matplotlib.pyplot as plt

yolov8_detect_npy = -np.ones((7000, 9))  # imgID,x1,y1,x2,y2,conf,label,gtID(iou>0.5),IOU
pycocotools_detect_npy = -np.ones((7000, 9)) # imgID,x1,y1,x2,y2,conf,label,gtID(iou>0.5),IOU
nanodet_detect_npy = -np.ones((7000, 9)) # imgID,x1,y1,x2,y2,conf,label,gtID(iou>0.5),IOU
gt_npy = -np.ones((7000, 7)) # imgID,x1,y1,x2,y2,conf,label

names = ['face', 'hand', 'cigarette', 'cellphone']

eps=1e-16

def smooth(y, f=0.05):
    # Box filter of fraction f
    nf = round(len(y) * f * 2) // 2 + 1  # number of filter elements (must be odd)
    p = np.ones(nf // 2)  # ones padding
    yp = np.concatenate((p * y[0], y, p * y[-1]), 0)  # y padded
    return np.convolve(yp, np.ones(nf) / nf, mode='valid')  # y-smoothed

def plot_mc_curve(px, py, save_dir=Path('mc_curve.png'), names=(), xlabel='Confidence', ylabel='Metric'):
    # Metric-confidence curve
    fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)

    if 0 < len(names) < 21:  # display per-class legend if < 21 classes
        for i, y in enumerate(py):
            ax.plot(px, y, linewidth=1, label=f'{names[i]}')  # plot(confidence, metric)
    else:
        ax.plot(px, py.T, linewidth=1, color='grey')  # plot(confidence, metric)

    y = smooth(py.mean(0), 0.05)
    ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}')
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    # 设置x刻度(是刻度不是刻度值)
    xticks_label = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    ax.set_xticks(xticks_label)
    yticks_label = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    ax.set_yticks(yticks_label)
    ax.legend(bbox_to_anchor=(1.04, 1), loc='upper left')
    ax.set_title(f'{ylabel}-Confidence Curve')
    fig.savefig(save_dir, dpi=250)
    plt.close(fig)

# 将yolov8保存的npy中的单张图局部gtID映射为全局gtID
def correctGtID():
    val_json = "C:/Users/Administrator/Desktop/1/val.json"
    imgID_gtID_dict = {}
    with open(val_json, "r", encoding="utf-8") as f:
        json_content = json.load(f)
        for ann in json_content["annotations"]:
            if ann["image_id"] not in imgID_gtID_dict.keys():
                imgID_gtID_dict[ann["image_id"]] = []
            imgID_gtID_dict[ann["image_id"]].append(ann["id"])
    # print(imgID_gtID_dict)

    yolov8_npy_file = "C:/Users/Administrator/Desktop/1/yolov8_detect_npy.npy"
    yolov8_result = np.load(yolov8_npy_file)
    yolov8_index = np.where(yolov8_result[:, 0] == -1.0)[0]
    yolov8_filtered_boxs = np.delete(yolov8_result, yolov8_index, axis=0)
    yolov8_result_final = np.array(sorted(yolov8_filtered_boxs, key=lambda x: x[0]))
    print("111111")
    for i, obj in enumerate(yolov8_result_final):
        imgID = obj[0]
        gtIDs = imgID_gtID_dict[imgID]
        lim_gtID = obj[7]
        if (int(lim_gtID) >= len(gtIDs)) or (int(lim_gtID) == -1):
            print("imgID:{} gtIDs:{} lim_gtID:{}".format(imgID, gtIDs, lim_gtID))
            continue
        global_gtID = gtIDs[int(lim_gtID)]
        yolov8_result_final[i][7] = global_gtID
    np.save("C:/Users/Administrator/Desktop/1/yolov8_detect_npy_new.npy", yolov8_result_final)


def readNpy():
    save_dir = "./tools/save/"
    #######################################GT#######################################
    val_json = "./tools/save/val.json"
    gt_num = 0
    gt_class_num_list = []
    with open(val_json, "r", encoding="utf-8") as f:
        json_content = json.load(f)
        gt_num = len(json_content["annotations"])
        print("gt_num:{}".format(gt_num))
        for i, ann in enumerate(json_content["annotations"]):
            gt_npy[i][0] = ann["image_id"]
            gt_npy[i][1] = ann["bbox"][0]
            gt_npy[i][2] = ann["bbox"][1]
            gt_npy[i][3] = ann["bbox"][0] + ann["bbox"][2]
            gt_npy[i][4] = ann["bbox"][1] + ann["bbox"][3]
            gt_npy[i][5] = 1.0
            gt_npy[i][6] = ann["category_id"]
    for c in range(4):
        c += 1
        gt_npy_singleC = gt_npy[np.where(gt_npy[:, 6] == c)[0]]
        gt_class_num = len(gt_npy_singleC)
        gt_class_num_list.append(gt_class_num)

    #######################################yolov8#######################################
    yolov8_npy_file = "./tools/save/yolov8_detect_npy_new.npy"
    yolov8_result = np.load(yolov8_npy_file)
    yolov8_index = np.where(yolov8_result[:, 0] == -1.0)[0]
    yolov8_filtered_boxs = np.delete(yolov8_result, yolov8_index, axis=0)
    yolov8_result_final = np.array(sorted(yolov8_filtered_boxs, key=lambda x: x[1]))
    yolov8_result_final = np.array(sorted(yolov8_result_final, key=lambda x: x[0]))
    # 保留小数位数
    for i, obj in enumerate(yolov8_result_final):
        yolov8_result_final[i][5] = round(yolov8_result_final[i][5], 4)
        yolov8_result_final[i][8] = round(yolov8_result_final[i][8], 4)
    # print(len(yolov8_result_final))
    # 计算AP
    detect_nums_yolov8 = yolov8_result_final.shape[0]
    tps_yolov8 = len(np.where(yolov8_result_final[:, -2] != -1.0)[0])
    fps_yolov8 = len(np.where(yolov8_result_final[:, -2] == -1.0)[0])
    precision_tra_yolov8 = tps_yolov8 / detect_nums_yolov8
    print("yolov8 detect_nums:{} tps:{} fps:{} precision_tra_yolov8:{}".format(detect_nums_yolov8, tps_yolov8, fps_yolov8, precision_tra_yolov8))
    ap_all_yolov8 = 0
    for c1 in range(len(names)):
        c1 += 1
        yolov8_result_final_singleC = yolov8_result_final[np.where(yolov8_result_final[:, 6] == c1)[0]]
        detect_nums = len(yolov8_result_final_singleC)
        tp_num = len(np.where(yolov8_result_final_singleC[:, -2] != -1)[0])
        fp_num = len(np.where(yolov8_result_final_singleC[:, -2] == -1)[0])
        precision_tra = tp_num / detect_nums
        print("yolov8 class:{} detect_nums:{} tp_num:{} fp_num:{} precision_tra:{}".format(c1, detect_nums, tp_num, fp_num, precision_tra))

        # 按照yolov8代码计算ap
        yolov8_result_final_singleC = np.array(sorted(yolov8_result_final_singleC, key=lambda x: x[5],reverse=True))
        tps = np.array([1 if (gtID!=-1) else 0 for gtID in yolov8_result_final_singleC[:, 7]])
        fpc = (1 - tps).cumsum(0)
        tpc = tps.cumsum(0)
        # Recall
        recall = tpc / (gt_class_num_list[c1-1] + eps)  # recall curve
        # Precision
        precision = tpc / (tpc + fpc)  # precision curve
        mrec = np.concatenate(([0.0], recall, [1.0]))
        mpre = np.concatenate(([1.0], precision, [0.0]))

        # Compute the precision envelope
        mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
        x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
        ap = np.trapz(np.interp(x, mrec, mpre), x)
        ap_all_yolov8 += ap
        print("yolov8 class:{} ap:{}".format(c1, ap))
    ap_all_yolov8 = ap_all_yolov8 / 4
    print("yolov8 ap_all:{}".format(ap_all_yolov8))

    ###################################pycocotools###########################################
    pycocotools_npy_file = "./tools/save/yolov8_detect_npy.npy"
    pycocotools_result = np.load(pycocotools_npy_file)
    pycocotools_index = np.where(pycocotools_result[:, 0] == -1.0)[0]
    pycocotools_filtered_boxs = np.delete(pycocotools_result, pycocotools_index, axis=0)
    pycocotools_result_final = np.array(sorted(pycocotools_filtered_boxs, key=lambda x: x[1]))
    pycocotools_result_final = np.array(sorted(pycocotools_result_final, key=lambda x: x[0]))
    for j, obj in enumerate(pycocotools_result_final):
        pycocotools_result_final[j][5] = round(pycocotools_result_final[j][5], 4)
        pycocotools_result_final[j][8] = round(pycocotools_result_final[j][8], 4)
    # print(len(pycocotools_result_final))
    # 计算AP
    detect_nums_pycocotools = pycocotools_result_final.shape[0]
    tps_pycocotools = len(np.where(pycocotools_result_final[:, -2] != -1.0)[0])
    fps_pycocotools = len(np.where(pycocotools_result_final[:, -2] == -1.0)[0])
    precision_tra_pycocotools = tps_pycocotools / detect_nums_pycocotools
    print("pycocotools detect_nums:{} tps:{} fps:{} precision_tra_pycocotools:{}".format(detect_nums_pycocotools, tps_pycocotools, fps_pycocotools, precision_tra_pycocotools))
    ap_all_pycocotools = 0
    p, r = np.zeros((4, 1000)), np.zeros((4, 1000))
    p1, r1 = np.ones((4, 1000)), np.zeros((4, 1000))
    for c2 in range(len(names)):
        c2 += 1
        pycocotools_result_final_singleC = pycocotools_result_final[np.where(pycocotools_result_final[:, 6] == c2)[0]]
        detect_nums = len(pycocotools_result_final_singleC)
        tps = len(np.where(pycocotools_result_final_singleC[:, -2] != -1)[0])
        fps = len(np.where(pycocotools_result_final_singleC[:, -2] == -1)[0])
        precision_tra = tps / detect_nums
        print("pycocotools class:{} detect_nums:{} tps:{} fps:{} precision:{}".format(c2, detect_nums, tps, fps, precision_tra))

        # 按照pycocotools代码计算ap
        tp = np.array([1 if (gtID != -1) else 0 for gtID in pycocotools_result_final_singleC[:, 7]])
        inds = np.argsort(-pycocotools_result_final_singleC[:, 5], kind='mergesort')
        tp = tp[inds]
        fpc = (1 - tp).cumsum(0)
        tpc = tp.cumsum(0)
        rc = tpc / gt_class_num_list[c2 - 1]
        pr = tpc / (fpc + tpc + np.spacing(1))
        q = np.zeros((101,))
        for i in range(detect_nums - 1, 0, -1):
            if pr[i] > pr[i - 1]:
                pr[i - 1] = pr[i]

        recThrs = np.linspace(0, 1, 101)  # 101-point interp (COCO)
        inds = np.searchsorted(rc, recThrs, side='left')
        try:
            for ri, pi in enumerate(inds):
                q[ri] = pr[pi]
        except:
            pass
        q_array = np.array(q)
        mean_s = np.mean(q_array[q_array > -1])
        print("pycocotools class:{} ap:{}".format(c2, mean_s))
        ap_all_pycocotools += mean_s

        # yolov8的ap计算方式
        mrec = np.concatenate(([0.0], rc, [1.0]))
        mpre = np.concatenate(([1.0], pr, [0.0]))
        # Compute the precision envelope
        mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))  # 使得低Recall的Precision不低于比它更高的Recall的Precision。
        x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
        temp = np.interp(x, mrec, mpre)  # 基本确定，yolov8推算的ap与pycocotools推算的ap差异就在此处，目前来看pycocotools的推算方式更合理
        mean_t = np.mean(temp)
        ap1 = np.trapz(np.interp(x, mrec, mpre), x)

        # 绘制PR曲线 两种曲线绘制方式视觉效果差不多
        # yolov8方式
        px = np.linspace(0, 1, 1000)
        i = np.argsort(-pycocotools_result_final_singleC[:, 5])
        conf = pycocotools_result_final_singleC[i, 5]
        p[c2-1] = np.interp(-px, -conf, pr, left=1)  # p at pr_score
        r[c2-1] = np.interp(-px, -conf, rc, left=0)  # negative x, xp because xp decreases
        # pycocotools方式
        conf1 = np.flip(conf)
        pr1 = np.flip(pr)
        rc1 = np.flip(rc)
        inds1 = np.searchsorted(conf1, px, side='left')
        try:
            for ri, pi in enumerate(inds1):
                p1[c2 - 1][ri] = pr1[pi]  # p at pr_score
                r1[c2 - 1][ri] = rc1[pi]  # negative x, xp because xp decreases
        except:
            pass

    ######################PR曲线-yolov8########################
    # P曲线
    plot_mc_curve(px, p, save_dir + 'P_curve.png', names, ylabel='Precision')
    # R曲线
    plot_mc_curve(px, r, save_dir + 'R_curve.png', names, ylabel='Recall')
    # ######################PR曲线-pycocotools########################
    # # P曲线
    # plot_mc_curve(px, p1, save_dir + 'P1_curve.png', names, ylabel='Precision')
    # # R曲线
    # plot_mc_curve(px, r1, save_dir + 'R1_curve.png', names, ylabel='Recall')

    ap_all_pycocotools = ap_all_pycocotools / 4
    print("pycocotools ap_all:{}".format(ap_all_pycocotools))

if __name__ == "__main__":
    readNpy()
    # correctGtID()